import random


import torch
import numpy as np
torch.set_default_tensor_type(torch.DoubleTensor)
import math

class Trainset_random_generated(object):
    def __init__(self, Batch, strating_snap, height, width):
        self.Batch = Batch
        self.strating_snap = strating_snap
        self.height = height
        self.width = width



        ##
    def __call__(self):
        torch.manual_seed(0)
        x=torch.rand(self.Batch, self.strating_snap, self.height, self.width)
        self.snapTrain= ((-1-1) * x+1)*0.9
        return self.snapTrain
class Trainset_random_generated0(object):
    def __init__(self, Batch, strating_snap, height, width):
        self.Batch = Batch
        self.strating_snap = strating_snap
        self.height = height
        self.width = width



        ##
    def __call__(self):
        torch.manual_seed(888)
        x = torch.rand(self.Batch, self.strating_snap, self.height, self.width)
        self.snapTrain= ((-1-1) * x+1)*0.9
        return self.snapTrain
class Trainset_PDE_generated(object):
    def __init__(self,strating_snap, height, width, epsilon,domain):
        self.strating_snap = strating_snap
        self.height = height
        self.width = width

        self.lbx = domain[0]
        self.ubx = domain[1]
        self.lby = domain[2]
        self.uby = domain[3]
        self.epsilon = epsilon

        self.xx = torch.linspace(self.lbx,self.ubx,steps = self.width+1)
        self.yy = torch.linspace(self.lby, self.uby, steps = self.height+1)
        ##
    def __call__(self):
        X, Y = torch.meshgrid(self.xx, self.yy)
        Z = torch.tanh((0.35 - torch.sqrt((X-0.5)*(X-0.5) + (Y-0.5) * (Y-0.5))) / torch.tensor(2*self.epsilon))
        self.snapTrain = torch.unsqueeze(torch.unsqueeze(Z, 0), 0)

        return self.snapTrain

class Trainset_noise_generated(object):
    def __init__(self, strating_snap, height, width, epsilon,domain=(-0.5, 0.5, -0.5, 0.5)):
        self.strating_snap = strating_snap
        self.height = height
        self.width = width
        self.lbx = domain[0]
        self.ubx = domain[1]
        self.lby = domain[2]
        self.uby = domain[3]
        self.epsilon = epsilon

        self.xx = np.linspace(self.lbx, self.ubx, np.int(self.width))
        self.yy = np.linspace(self.lby, self.uby, np.int(self.height))

    def __call__(self):
        X, Y = np.meshgrid(self.xx, self.yy)
        Z = np.tanh(0.4-np.sqrt(X*X+Y*Y))/(np.sqrt(2)*self.epsilon)
        ZZ = []
        noise = 0
        for i in range(self.strating_snap):
            Z = Z + 0.01 * noise
            noise = np.random.rand(self.height, self.width)
            ZZ.append(Z)
        ZZ = np.array(ZZ)
        self.snapTrain = torch.from_numpy(ZZ)
        self.snapTrain = torch.unsqueeze(self.snapTrain, 1)
        return self.snapTrain



# if __name__ == '__main__':
